#!/usr/bin/perl
# 2013/05/08 Make association map for ChIA-PET


use strict;
use warnings;
use IO::File;
use Getopt::Std;
use File::Basename;
use Carp qw(croak);
$| = 0;


if(@ARGV != 8 or $ARGV[0] eq '--help'){
	die "Usage : $0 -i [map file] -o [output association file(prefix for format2)] -r [resolution] -f [format type: 1:aso, 2:matrix]\n";
}

my %opt;
getopts("i:r:o:f:", \%opt);
my $FILE_map = $opt{i};
my $Resolution = $opt{r};
my $FILE_out = $opt{o};
my $OUT_FORMAT = $opt{f};

my %CHROMOSOME = (
	'chr1' => 'I',
	'chr2' => 'II',
	'chr3' => 'III',
	'I' => 'I',
	'II' => 'II',
	'III' => 'III'
);


my %Len_chr = (
	'I'	=>	5579133,
	'II'	=>	4539804,
	'III'	=>	2452883,
);



#---------------------------------------
# correct information
#---------------------------------------
my %data;
my $fh_map = IO::File->new($FILE_map) or die "cannot open $FILE_map: $!";
while($_ = $fh_map->getline()){
	s/\r?\n//;
	my ($id, $chr1, $loc1, $direction1, $mapQ1, $hinID1, $hinLoc1, $chr2, $loc2, $direction2, $mapQ2, $hinID2, $hinLoc2, $type, $num, $locations) = split /\t/;

	my $cate1 = int($loc1 / $Resolution) * $Resolution + 1;
	my $cate1_end = $cate1 + $Resolution - 1;
	my $cate2 = int($loc2 / $Resolution) * $Resolution + 1;
	my $cate2_end = $cate2 + $Resolution - 1;

	if(exists $CHROMOSOME{$chr1}){
		$chr1 = $CHROMOSOME{$chr1};
	}else{
		next;
	}

	if(exists $CHROMOSOME{$chr2}){
		$chr2 = $CHROMOSOME{$chr2};
	}else{
		next;
	}

	my $bin1 = "$chr1:$cate1-$cate1_end";
	my $bin2 = "$chr2:$cate2-$cate2_end";


	if($chr1 eq $chr2){
		if($loc1 < $loc2){
			$data{$bin1}{$bin2}++;
		}else{
			$data{$bin2}{$bin1}++;
		}
	}else{
		if($chr1 eq 'I'){
			$data{$bin1}{$bin2}++;
		}elsif($chr1 eq 'II' and $chr2 eq 'III'){
			$data{$bin1}{$bin2}++;
		}else{
			$data{$bin2}{$bin1}++;
		}
	}
}
$fh_map->close();


#---------------------------------------
# register bins
#---------------------------------------
my @binLists;
foreach my $chr(qw(I II III)){
	for(my $i = 1; $i < $Len_chr{$chr}; $i += $Resolution){
		my $end = $i + $Resolution - 1;
		push @binLists, "$chr:$i-$end";
	}
}


#---------------------------------------
# output
#---------------------------------------
if($OUT_FORMAT == 1){
	my $fh_out = IO::File->new($FILE_out, 'w') or die "cannot write $FILE_out: $!";
	for(my $i = 0; $i < @binLists; $i++){
		my $bin1 = $binLists[$i];
		for(my $j = 0; $j < @binLists; $j++){
			my $bin2 = $binLists[$j];
			my $value = 0;
			if($i < $j){
				$value = exists $data{$bin1}{$bin2} ? $data{$bin1}{$bin2} : 0;
			}else{
				$value = exists $data{$bin2}{$bin1} ? $data{$bin2}{$bin1} : 0;
			}
			$fh_out->print("$bin1\t$bin2\t$value\n");
		}
	}
	$fh_out->close();
}elsif($OUT_FORMAT == 2){
	my $FILE_out1 = $FILE_out . "_chr1.matrix";
	my $FILE_out2 = $FILE_out . "_chr2.matrix";
	my $FILE_out3 = $FILE_out . "_chr3.matrix";
	my $fh_out1 = IO::File->new($FILE_out1, 'w') or die "cannot write $FILE_out1: $!";
	my $fh_out2 = IO::File->new($FILE_out2, 'w') or die "cannot write $FILE_out2: $!";
	my $fh_out3 = IO::File->new($FILE_out3, 'w') or die "cannot write $FILE_out3: $!";
	my %FH_OUT = (
		'I' => $fh_out1,
		'II' => $fh_out2,
		'III' => $fh_out3
	);
	foreach my $chr_target (qw(I II III)){
		for(my $i = 0; $i < @binLists; $i++){
			my $bin1 = $binLists[$i];
			my ($chr1, $loc1) = split /:/, $bin1;
			if($chr1 ne $chr_target){
				next;
			}
			my ($start1, $end1) = split /\-/, $loc1;
			$FH_OUT{$chr_target}->print("$chr1\t$start1\t$end1");
			for(my $j = 0; $j < @binLists; $j++){
				my $bin2 = $binLists[$j];
				my ($chr2, $loc2) = split /:/, $bin2;
				if($chr2 ne $chr_target){
					next;
				}
				my $value = 0;
				if($i < $j){
					$value = exists $data{$bin1}{$bin2} ? $data{$bin1}{$bin2} : 0;
				}else{
					$value = exists $data{$bin2}{$bin1} ? $data{$bin2}{$bin1} : 0;
				}
				$FH_OUT{$chr_target}->print("\t$value");
			}
			$FH_OUT{$chr_target}->print("\n");
		}
	}
	$fh_out1->close();
	$fh_out2->close();
	$fh_out3->close();
	system("gzip $FILE_out1") == 0 or die "cannot commpress\n";
	system("gzip $FILE_out2") == 0 or die "cannot commpress\n";
	system("gzip $FILE_out3") == 0 or die "cannot commpress\n";
}


